import numpy as np
import cv2

# 根据分类数产生对应颜色
def bit_get(val, idx):
    """Gets the bit value.
    Args:
      val: Input value, int or numpy int array.
      idx: Which bit of the input val.
    Returns:
      The "idx"-th bit of input val.
    """
    return (val >> idx) & 1

def create_pascal_label_colormap(class_num):
    """Creates a label colormap used in PASCAL VOC segmentation benchmark.
    Returns:
      A colormap for visualizing segmentation results.
    """
    colormap = np.zeros((class_num, 3), dtype=int)
    ind = np.arange(class_num, dtype=int)

    for shift in reversed(range(8)):
        for channel in range(3):
            colormap[:, channel] |= bit_get(ind, channel) << shift
        ind >>= 3

    return colormap

# 用于显示语义分割结果
def to_color_img(img, color_list):

    score_i = img.argmax(dim=1)

    score_i = score_i.squeeze(0)

    score_i = score_i.cpu().numpy()
    color_img = color_list[score_i].astype(np.uint8)
    return color_img
    
def label_img_to_color(img):
    label_to_color = {
        0: [128, 64, 128],
        1: [255, 255, 255],
        2: [ 70, 70, 70],
        3: [102,102,156],
        4: [190,153,153],
        5: [153,153,153],
        6: [250,170, 30],
        7: [220,220,  0],
        8: [107,142, 35],
        9: [152,251,152],
        10: [ 70,130,180],
        11: [220, 20, 60],
        12: [255,  0,  0],
        13: [  0,  0,142],
        14: [  0,  0, 70],
        15: [  0, 60,100],
        16: [  0, 80,100],
        17: [  0,  0,230],
        18: [119, 11, 32],
        19: [81,  0, 81]
        }

    _, _, img_height, img_width = img.shape
    # img_height, img_width = img.shape

    img_color = np.zeros((img_height, img_width, 3), dtype=np.uint8)
    score_i = img.argmax(dim=1).squeeze(0).cpu().numpy()
    # score_i = img

    for row in range(img_height):
        for col in range(img_width):
            label = score_i[row, col]

            img_color[row, col] = np.array(label_to_color[label])
            print(label_to_color[label])
            break
    return img_color
if __name__ == "__main__":
    img = "/workspace/wzj/dataset/cityscape/gtFine/train/aachen/aachen_000000_000019_gtFine_labelTrainIds.png"
    
    
    img = cv2.imread(img, 0)
    print(img.shape)
    # img = cv2.resize(img, (1024, 512), interpolation=cv2.INTER_NEAREST)
    
    # img = label_img_to_color(img)
    # cv2.imshow("img", img)
    # cv2.waitKey()
    # cv2.destroyAllWindows()